#coding:utf8
'''
Created on 2016年5月23日

@author: zhangjinrui
'''

import psutil
import socket
import json
import time

class Collector(object):


    def __init__(self, items_count, get_aimed_addrs = False):
        self.get_aimed_addrs = get_aimed_addrs
        self.items_count = items_count
        self.hostname = socket.gethostname()
        self.ip = socket.gethostbyname(self.hostname)
    
    def start(self):
        #创建文件,文件名以ip命名
        path = '../data/' + self.ip + '.json'
        with open(path,'w') as f:
            disk_info = psutil.disk_io_counters(False)
            net_info = psutil.net_io_counters(False)
            
            while(self.items_count > 0 or self.items_count == -1):
                #获取CPU信息
                cpu_each_percent = psutil.cpu_percent(0.1, True)
                cpu_all_percent = psutil.cpu_percent(0.1, False)
                
                #获取内存信息
                memory_info = psutil.virtual_memory()
                memory_percent = memory_info.percent
                
                #获取磁盘信息
                last_disk_info = disk_info
                disk_info = psutil.disk_io_counters(False)
                #磁盘速度用KB/S表示
                disk_read_speed = 1.0 * (disk_info.read_bytes - last_disk_info.read_bytes) / 1024
                disk_write_speed = 1.0 * (disk_info.write_bytes - last_disk_info.write_bytes) / 1024
                
                #获取网络信息
                last_net_info = net_info
                net_info = psutil.net_io_counters(False)
                net_rec_speed = 1.0 * (net_info.bytes_recv - last_net_info.bytes_recv) / 1024
                net_send_speed = 1.0 * (net_info.bytes_sent - last_net_info.bytes_sent) / 1024
                
                #获取网络发出目标列表
                aimed_addrs = []
                if self.get_aimed_addrs == True:
                    aimed_addrs_dict = {}
                    net_conns_info = psutil.net_connections('inet4');
                    for addr_proc in net_conns_info:
                        addr = addr_proc.raddr
                        if(len(addr) > 0):
                            ip = addr[0]
                        else:
                            continue
                        
                        #过滤掉和本地的通讯
                        if ip == '127.0.0.1' or ip == '0.0.0.0' or ip == 'localhost':
                            continue
                        
                        if not aimed_addrs_dict.has_key(ip):
                            aimed_addrs.append(ip)
                            aimed_addrs_dict[ip] = 1
                
                #将数据打包成json
                obj = {}
                obj['timestamp'] = time.time()
                obj['hostname'] = self.hostname
                obj['ip'] = self.ip
                obj['cpu_each_percent'] = cpu_each_percent
                obj['cpu_all_percent'] = cpu_all_percent
                obj['memory_percent'] = memory_percent
                obj['disk_read_speed'] = disk_read_speed
                obj['disk_write_speed'] = disk_write_speed
                obj['net_rec_speed'] = net_rec_speed
                obj['net_send_speed'] = net_send_speed
                obj['aimed_addrs'] = aimed_addrs
                json_string = json.dumps(obj)
                
                f.write(json_string + '\n')
                
                if(self.items_count != -1):
                    self.items_count -= 1
                
                time.sleep(0.8)
        
    